#!/usr/bin/env python
# -- coding: utf-8 --
"""
Copyright (c) 2018. All rights reserved.
Created by C. L. Wang on 2018/4/18

参考:
NumPy FutureWarning
https://stackoverflow.com/questions/48340392/futurewarning-conversion-of-the-second-argument-of-issubdtype-from-float-to
"""

import os

import click
import numpy as np

from infers.simple_mnist_infer import SimpleMnistInfer
from loaders.simple_mnist_dl import SimpleMnistDL
from models.simple_mnist_model import SimpleMnistModel
from trainers.simple_mnist_trainer import SimpleMnistTrainer
from utils.config_utils import get_train_args, process_config
from utils.logger import log


@click.group()
@click.option('-v', '--verbose', count=True)
@click.pass_context
def cli(ctx, verbose):
    ctx.obj["VERBOSE"] = verbose


@cli.command(help='创建一个实例.')
@click.option('-n', '--name', default=None, help='实例名称.')
def create(name):
    fname = '{}.py'.format(name)
    files = []
    files.append(os.path.join('configs', '{}_config.json'.format(name)))
    files.append(os.path.join('loaders', '{}_dl.py'.format(name)))
    files.append(os.path.join('models', '{}_model.py'.format(name)))
    files.append(os.path.join('infers', '{}_infer.py'.format(name)))

    for x in files:
        if not os.path.exists(x):
            with open(x,'w') as fp:
                fp.write('{}')

@cli.command(help='训练一个实例.')
@click.option('-n', '--name', default=None, help='实例名称.')
def train(name):
    log.info('解析配置...')
    config = process_config('configs/{}_config.json'.format(name))

    np.random.seed(47)  # 固定随机数

    log.info('加载数据...')
    dl = SimpleMnistDL(config=config)

    log.info('构造网络...')
    model = SimpleMnistModel(config=config)

    log.info('训练网络...')
    trainer = SimpleMnistTrainer(
        model=model.model,
        data=[dl.get_train_data(), dl.get_test_data()],
        config=config)
    trainer.train()
    
    log.info('训练完成...')

@cli.command(help='测试一个实例.')
@click.option('-n', '--name', default=None, help='实例名称.')
@click.option('-m', '--model', default=None, help='实例名称.')
def tests(name, model):
    log.info('解析配置...')
    config = process_config('configs/{}_config.json'.format(name))

    log.info('加载数据...')
    
    dl = SimpleMnistDL()
    test_data = np.expand_dims(dl.get_test_data()[0][0], axis=0)
    test_label = np.argmax(dl.get_test_data()[1][0])

    log.info('预测数据...')

    model = "simple_mnist.weights.01-1.61.hdf5" if model is None else model
    model = model.replace('.hdf5', '') + '.hdf5'
    infer = SimpleMnistInfer(model, config)
    infer_label = np.argmax(infer.predict(test_data))
    
    log.info('真实Label: %s, 预测Label: %s' % (test_label, infer_label))
    log.info('预测完成...') 


def entry():
    cli(obj={})


if __name__ == "__main__":
    entry()
